# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import ast


def argsparser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--det_model_dir",
        type=str,
        default=None,
        help=(
            "Directory include:'model.pdiparams', 'model.pdmodel', "
            "'infer_cfg.yml', created by tools/export_model.py."
        ),
        required=True,
    )
    parser.add_argument(
        "--keypoint_model_dir",
        type=str,
        default=None,
        help=(
            "Directory include:'model.pdiparams', 'model.pdmodel', "
            "'infer_cfg.yml', created by tools/export_model.py."
        ),
        required=True,
    )
    parser.add_argument("--image_file", type=str, default=None, help="Path of image file.")
    parser.add_argument(
        "--image_dir",
        type=str,
        default=None,
        help="Dir of image file, `image_file` has a higher priority.",
    )
    parser.add_argument(
        "--keypoint_batch_size",
        type=int,
        default=8,
        help=(
            "batch_size for keypoint inference. In detection-keypoint unit"
            "inference, the batch size in detection is 1. Then collate det "
            "result in batch for keypoint inference."
        ),
    )
    parser.add_argument(
        "--video_file",
        type=str,
        default=None,
        help="Path of video file, `video_file` or `camera_id` has a highest priority.",
    )
    parser.add_argument("--camera_id", type=int, default=-1, help="device id of camera to predict.")
    parser.add_argument("--det_threshold", type=float, default=0.5, help="Threshold of score.")
    parser.add_argument("--keypoint_threshold", type=float, default=0.5, help="Threshold of score.")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output",
        help="Directory of output visualization files.",
    )
    parser.add_argument(
        "--run_mode",
        type=str,
        default="paddle",
        help="mode of running(paddle/trt_fp32/trt_fp16/trt_int8)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        help="Choose the device you want to run, it can be: CPU/GPU/XPU, default is CPU.",
    )
    parser.add_argument(
        "--run_benchmark",
        type=ast.literal_eval,
        default=False,
        help="Whether to predict a image_file repeatedly for benchmark",
    )
    parser.add_argument(
        "--enable_mkldnn",
        type=ast.literal_eval,
        default=False,
        help="Whether use mkldnn with CPU.",
    )
    parser.add_argument("--cpu_threads", type=int, default=1, help="Num of threads with CPU.")
    parser.add_argument("--trt_min_shape", type=int, default=1, help="min_shape for TensorRT.")
    parser.add_argument("--trt_max_shape", type=int, default=1280, help="max_shape for TensorRT.")
    parser.add_argument("--trt_opt_shape", type=int, default=640, help="opt_shape for TensorRT.")
    parser.add_argument(
        "--trt_calib_mode",
        type=bool,
        default=False,
        help="If the model is produced by TRT offline quantitative " "calibration, trt_calib_mode need to set True.",
    )
    parser.add_argument(
        "--use_dark",
        type=ast.literal_eval,
        default=True,
        help="whether to use darkpose to get better keypoint position predict ",
    )
    parser.add_argument(
        "--save_res",
        type=bool,
        default=False,
        help=(
            "whether to save predict results to json file"
            "1) store_res: a list of image_data"
            "2) image_data: [imageid, rects, [keypoints, scores]]"
            "3) rects: list of rect [xmin, ymin, xmax, ymax]"
            "4) keypoints: 17(joint numbers)*[x, y, conf], total 51 data in list"
            "5) scores: mean of all joint conf"
        ),
    )
    parser.add_argument(
        "--smooth",
        type=ast.literal_eval,
        default=False,
        help="smoothing keypoints for each frame, new incoming keypoints will be more stable.",
    )
    parser.add_argument(
        "--filter_type",
        type=str,
        default="OneEuro",
        help="when set --smooth True, choose filter type you want to use, it can be [OneEuro] or [EMA].",
    )
    return parser
